package com.pt.util;

import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.List;
import java.util.HashMap;
import com.pt.entity.Peer;

public class BencodeUtil {
    
    public static byte[] encodeTorrent(Map<String, Object> data) {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        try {
            encodeObject(data, baos);
            return baos.toByteArray();
        } catch (Exception e) {
            throw new RuntimeException("Bencode编码失败", e);
        }
    }
    
    private static void encodeObject(Object obj, ByteArrayOutputStream baos) throws Exception {
        if (obj instanceof String) {
            encodeString((String) obj, baos);
        } else if (obj instanceof Integer) {
            encodeInteger((Integer) obj, baos);
        } else if (obj instanceof Long) {
            encodeLong((Long) obj, baos);
        } else if (obj instanceof List) {
            encodeList((List<?>) obj, baos);
        } else if (obj instanceof Map) {
            encodeMap((Map<?, ?>) obj, baos);
        } else if (obj instanceof byte[]) {
            encodeBytes((byte[]) obj, baos);
        } else {
            throw new IllegalArgumentException("不支持的类型: " + obj.getClass());
        }
    }
    
    private static void encodeString(String str, ByteArrayOutputStream baos) throws Exception {
        byte[] bytes = str.getBytes(StandardCharsets.ISO_8859_1);
        baos.write((bytes.length + ":").getBytes(StandardCharsets.ISO_8859_1));
        baos.write(bytes);
    }
    
    private static void encodeInteger(Integer num, ByteArrayOutputStream baos) throws Exception {
        baos.write(("i" + num + "e").getBytes(StandardCharsets.ISO_8859_1));
    }
    
    private static void encodeLong(Long num, ByteArrayOutputStream baos) throws Exception {
        baos.write(("i" + num + "e").getBytes(StandardCharsets.ISO_8859_1));
    }
    
    private static void encodeList(List<?> list, ByteArrayOutputStream baos) throws Exception {
        baos.write('l');
        for (Object item : list) {
            encodeObject(item, baos);
        }
        baos.write('e');
    }
    
    private static void encodeMap(Map<?, ?> map, ByteArrayOutputStream baos) throws Exception {
        baos.write('d');
        for (Map.Entry<?, ?> entry : map.entrySet()) {
            if (!(entry.getKey() instanceof String)) {
                throw new IllegalArgumentException("Map的键必须是String类型");
            }
            encodeString((String) entry.getKey(), baos);
            encodeObject(entry.getValue(), baos);
        }
        baos.write('e');
    }
    
    private static void encodeBytes(byte[] bytes, ByteArrayOutputStream baos) throws Exception {
        baos.write((bytes.length + ":").getBytes(StandardCharsets.ISO_8859_1));
        baos.write(bytes);
    }
    
    public static byte[] createErrorResponse(String message) {
        try {
            Map<String, Object> response = new HashMap<>();
            response.put("failure reason", message);
            return encodeTorrent(response);
        } catch (Exception e) {
            return ("d14:failure reason" + message.length() + ":" + message + "e").getBytes(StandardCharsets.ISO_8859_1);
        }
    }
    
    public static byte[] encodePeerResponse(List<Peer> peers, int seeders, int leechers) {
        try {
            Map<String, Object> response = new HashMap<>();
            response.put("interval", 1800);
            response.put("min interval", 900);
            response.put("complete", seeders);
            response.put("incomplete", leechers);
            
            ByteArrayOutputStream peerBytes = new ByteArrayOutputStream();
            for (Peer peer : peers) {
                String ip = peer.getIp();
                if ("0:0:0:0:0:0:0:1".equals(ip) || "::1".equals(ip)) {
                    ip = "127.0.0.1";
                }
                
                if (ip.contains(".")) {
                    String[] ipParts = ip.split("\\.");
                    if (ipParts.length == 4) {
                        for (String part : ipParts) {
                            int ipPart = Integer.parseInt(part);
                            peerBytes.write(ipPart & 0xFF);
                        }
                    }
                }
                
                int port = peer.getPort();
                peerBytes.write((port >> 8) & 0xFF);
                peerBytes.write(port & 0xFF);
            }
            
            response.put("peers", peerBytes.toByteArray());
            return encodeTorrent(response);
        } catch (Exception e) {
            throw new RuntimeException("编码Peer响应失败", e);
        }
    }
} 